import time
from mylib import UartCommend, Camera
import tf

# 初始化串口通信
uartcommend = UartCommend(uart_port=1, baudrate=115200)

# 初始化摄像头
camera = Camera()

# 模型和标签文件路径（SD卡根目录）
model_path = "/sd/model.tflite"
label_path = "/sd/label.txt"

# 加载模型和标签
print(f"Loading model: {model_path}")
tf_model = tf.load(model_path)

print(f"Loading labels: {label_path}")
with open(label_path, "r") as f:
    labels = [line.strip() for line in f]

def run_model(image):
    """
    在图像上运行模型并返回分类结果。
    """
    for obj in tf.classify(tf_model, image, min_scale=1.0, scale_mul=0.5, x_overlap=-1, y_overlap=-1):
        # 获取分类结果并排序
        sorted_list = sorted(zip(labels, obj.output()), key=lambda x: x[1], reverse=True)
        top_result = sorted_list[0]  # 获取最可信的结果
        return top_result  # 返回 (标签, 置信度)
    return None  # 如果没有分类结果

while True:
    # 获取图像和FPS
    img_roi, fps = camera.get_image()
    # print("FPS:", fps)

    # 处理接收到的串口数据
    while uartcommend.uart.any():
        data = uartcommend.get_uart_data()
        if data:
            if data[0] == 0x01:  # 如果数据是0x01，打开LED并运行模型
                camera.led.on()
                uartcommend.send_uart_data(0x26)
                # 跑模型
                result = run_model(img_roi)
                if result:
                    print(f"Detected: {result[0]}, Confidence: {result[1]:.2f}")
                    uartcommend.send_uart_data(0x66)  # 发送确认消息 向左撞击
                else:
                    print("No detection")
            elif data[0] == 0x02:  # 如果数据是0x02，关闭LED并运行模型
                camera.led.off()
                # 跑模型
                result = run_model(img_roi)
                if result:
                    print(f"Detected: {result[0]}, Confidence: {result[1]:.2f}")
                    uartcommend.send_uart_data(0x02)  # 发送确认消息 向右撞击
                else:
                    print("No detection")
